package websocket_handler

import (
	"fmt"
	"log"
	"net/http"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// WebSocketUpgrader 用于将 HTTP 连接升级为 WebSocket 连接
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		// 允许所有来源请求，可以根据需要自定义
		return true
	},
}

// ConnectionManager 管理所有 WebSocket 连接
type ConnectionManager struct {
	clients map[*websocket.Conn]bool // 存储所有活跃的 WebSocket 连接
	mu      sync.Mutex               // 保证对 clients 的并发访问安全
}

// NewConnectionManager 创建一个新的连接管理器
func NewConnectionManager() *ConnectionManager {
	return &ConnectionManager{
		clients: make(map[*websocket.Conn]bool),
	}
}

// AddClient 添加新的 WebSocket 连接
func (cm *ConnectionManager) AddClient(conn *websocket.Conn) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	cm.clients[conn] = true
}

// RemoveClient 移除 WebSocket 连接
func (cm *ConnectionManager) RemoveClient(conn *websocket.Conn) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	delete(cm.clients, conn)
}

// Broadcast 向所有客户端发送消息
func (cm *ConnectionManager) Broadcast(message string) {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	for conn := range cm.clients {
		fmt.Println("向客户端发送消息")
		fmt.Println(message)
		err := conn.WriteMessage(websocket.TextMessage, []byte(message))
		if err != nil {
			log.Println("Error writing message to client:", err)
			conn.Close()
			delete(cm.clients, conn) // 移除无效连接
		}
	}
}

// 全局连接管理器实例
var connectionManager = NewConnectionManager()

// HandleWebSocket WebSocket 连接处理函数
func HandleWebSocket(c *gin.Context) {
	sessionID := c.Param("sid")
	log.Printf("Received WebSocket connection for session ID: %s", sessionID)

	// 升级协议
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println("Failed to upgrade connection:", err)
		return
	}
	defer conn.Close()

	// 将连接加入到连接管理器
	connectionManager.AddClient(conn)
	defer connectionManager.RemoveClient(conn)

	for {
		// 读取客户端消息
		_, message, err := conn.ReadMessage()
		if err != nil {
			log.Printf("WebSocket connection closed for session %s: %v", sessionID, err)
			break
		}
		log.Printf("Received message for session %s: %s", sessionID, message)

		// 回复一条消息
		err = conn.WriteMessage(websocket.TextMessage, []byte("Hello from server!"))
		if err != nil {
			log.Println("Error writing message:", err)
			break
		}
	}
}

// BroadcastMessage 向所有客户端广播消息
func BroadcastMessage(message string) {
	connectionManager.Broadcast(message)
}
